import os, sys, zipfile
import json
import glob
import shutil
from tqdm import tqdm
from PIL import Image
import numpy as np
import pandas as pd
import cv2
import tifffile.tifffile as tifffile


class_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)
    




def generate_lables(json_path, save_label_path):
    # 打开标注文件
    data = json.load(open(json_path, 'r'))
    id_name_mapping = {i['id']: i['name'] for i in data['categories']}
    # 保存的路径
    ana_txt_save_path = save_label_path  

    for img in data['images']:
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        ana_txt_name = filename.split(".")[0] + ".txt"  # 对应的txt名字，与jpg一致
        print(ana_txt_name)
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                # import ipdb;ipdb.set_trace()
                box = convert((img_width, img_height), ann["bbox"])
                # import ipdb;ipdb.set_trace()
                cls_id = class_names.index(id_name_mapping[ann["category_id"]])
                f_txt.write("%s %s %s %s %s\n" % (cls_id, box[0], box[1], box[2], box[3]))
        f_txt.close()


def handle_hy():

    front_train_image = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_front/train_img/*.tif')
    front_train_mask = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_front/train_parsing/*.png')
    front_val_image = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_front/val_img/*.tif')
    front_val_mask = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_front/val_parsing/*.png')

    back_train_image = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_back/train_img/*.tif')
    back_train_mask = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_back/train_parsing/*.png')
    back_val_image = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_back/val_img/*.tif')
    back_val_mask = glob.glob(r'/home/workspace/data/data_v2/CIHP_body_back/val_parsing/*.png')


    for i in tqdm(front_train_image):
        name = i.split('/')[-1]
        destination_name = 'front_' + name

        mask_path = i.replace('train_img', 'train_parsing').replace('tif', 'png')
        if os.path.exists(mask_path):
            destination_mask = mask_path.split('/')[-1]
            destination_mask_name = 'front_' + destination_mask
            shutil.copy(i, '/home/workspace/data/yolou-hy/images/train/%s' % destination_name)
            shutil.copy(mask_path, '/home/workspace/data/yolou-hy/images/train/%s' % destination_mask_name)


    for i in tqdm(front_val_image):
        name = i.split('/')[-1]
        destination_name = 'front_' + name

        mask_path = i.replace('val_img', 'val_parsing').replace('tif', 'png')
        if os.path.exists(mask_path):
            destination_mask = mask_path.split('/')[-1]
            destination_mask_name = 'front_' + destination_mask
            shutil.copy(i, '/home/workspace/data/yolou-hy/images/val/%s' % destination_name)
            shutil.copy(mask_path, '/home/workspace/data/yolou-hy/images/val/%s' % destination_mask_name)




    for i in tqdm(back_train_image):
        name = i.split('/')[-1]
        destination_name = 'back_' + name

        mask_path = i.replace('train_img', 'train_parsing').replace('tif', 'png')
        if os.path.exists(mask_path):
            destination_mask = mask_path.split('/')[-1]
            destination_mask_name = 'back_' + destination_mask
            shutil.copy(i, '/home/workspace/data/yolou-hy/images/train/%s' % destination_name)
            shutil.copy(mask_path, '/home/workspace/data/yolou-hy/images/train/%s' % destination_mask_name)


    for i in tqdm(back_val_image):
        name = i.split('/')[-1]
        destination_name = 'back_' + name

        mask_path = i.replace('val_img', 'val_parsing').replace('tif', 'png')
        if os.path.exists(mask_path):
            destination_mask = mask_path.split('/')[-1]
            destination_mask_name = 'back_' + destination_mask
            shutil.copy(i, '/home/workspace/data/yolou-hy/images/val/%s' % destination_name)
            shutil.copy(mask_path, '/home/workspace/data/yolou-hy/images/val/%s' % destination_mask_name)




def generate_lables_hy(image_folder, save_label_path):
    # 打开标注文件
    # data = json.load(open(json_path, 'r'))
    # id_name_mapping = {i['id']: i['name'] for i in data['categories']}
    # 保存的路径
    for root, _, files in os.walk(image_folder):
        for i in tqdm(files):
            if i.endswith('png'):
                path = os.path.join(root, i)
                image = np.array(Image.open(path))
                image[image>1] = 1
                # df = pd.DataFrame(image)
                # ss = df.where(df=1)
                ss = np.where(image==1)


                ymax = np.max(ss[0])
                ymin = np.min(ss[0]) - 5

                xmax = np.max(ss[1]) + 5
                xmin = np.min(ss[1]) - 5


                # image = tifffile.imread(path.replace('png', 'tif'))
                # draw1 = cv2.rectangle(np.uint8(image) * 5, (xmin, ymin), (xmax, ymax), (0, 0, 255), 1)
                # cv2.imwrite("%s.png" % i, draw1)
                # cv2.imshow("test", draw1)

                width = xmax - xmin
                height = ymax - ymin

                ana_txt_save_path = save_label_path  



                img_width = image.shape[1]
                img_height = image.shape[0]
                ana_txt_name = i.split(".")[0] + ".txt"  # 对应的txt名字，与jpg一致
                # print(ana_txt_name)
                f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
                # import ipdb;ipdb.set_trace()
                box = [xmin, ymin, width, height]
                box = convert((img_width, img_height), box)
                # import ipdb;ipdb.set_trace()
                # cls_id = class_names.index(id_name_mapping[ann["category_id"]])
                f_txt.write("%s %s %s %s %s\n" % (0, box[0], box[1], box[2], box[3]))
                f_txt.close()



if __name__=='__main__':
    # generate_lables('/home/workspace/data/yolou-data/annotations/instances_train2017.json', '/home/workspace/data/yolou-data/labels/train2017')
    # generate_lables('/home/workspace/data/yolou-data/annotations/instances_val2017.json', '/home/workspace/data/yolou-data/labels/val2017')
    # handle_hy()
    generate_lables_hy('/home/workspace/data/yolou-hy/images/train/', '/home/workspace/data/yolou-hy/labels/train/')
    generate_lables_hy('/home/workspace/data/yolou-hy/images/val/', '/home/workspace/data/yolou-hy/labels/val/')